{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DKbZ5kitLyCX"
   },
   "source": [
    "##### Copyright 2020 The TensorFlow Authors.\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OAnLZy0vL0iF"
   },
   "outputs": [],
   "source": [
    "#@title ##### Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RalQobAv2IU4"
   },
   "source": [
    "# Single Column MNIST Classification\n",
    "\n",
    "In this colab, we experiment with classification of the traditional MNIST dataset, but in the presence of multimodal uncertainty.  We accomplish this by degrading the input signal:  Feeding the model only the central column from the input image.  By doing so, we cause certain output labels (e.g. 5, 8, 3) to have inputs which appear nearly identical.  We also demonstrate how to model this data using Variational Inference with mixture distributions for the posterior.  This colab mirrors one of the experiments presented in [1], in which we demonstrate that optimizing a Stratified-IWAE bound on the evidence returns a model which is better able to capture multimodal uncertainty than models which optimize a Stratified-ELBO bound.\n",
    "\n",
    "### References\n",
    "\n",
    "[1] Morningstar et al. (2020):  Automatic Differentiation Variational Inference with Mixtures ([arxiv preprint](https://arxiv.org/abs/2003.01687))\n",
    "\n",
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/probability/blob/master/tensorflow_probability/python/experimental/nn/examples/single_column_mnist.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/tensorflow/probability/blob/master/tensorflow_probability/python/experimental/nn/examples/single_column_mnist.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
    "  </td>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QV16_gj7EzIm"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import functools\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "import tensorflow.compat.v2 as tf\n",
    "import tensorflow_probability as tfp\n",
    "import tensorflow_datasets as tfds\n",
    "tf.enable_v2_behavior()\n",
    "\n",
    "from tensorflow_probability.python.experimental import nn as tfn\n",
    "\n",
    "tfb = tfp.bijectors\n",
    "tfd = tfp.distributions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "A1Cv8ACng85k"
   },
   "source": [
    "# Getting the data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JFTSraX4E8JM"
   },
   "outputs": [],
   "source": [
    "class Preprocess(object):\n",
    "  def __init__(self, info):\n",
    "    \"\"\"Initialize the preprocessing object.\"\"\"\n",
    "\n",
    "    self.ind = info.features['image'].shape[0] // 2\n",
    "\n",
    "  def __call__(self, inputs):\n",
    "    # Normalize and center the input image\n",
    "    x = 2*tf.cast(inputs['image'][..., self.ind, 0], tf.float32) / 255. -1.\n",
    "    y = tf.cast(inputs['label'], tf.int32)\n",
    "\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nw5KYALjIbXw"
   },
   "outputs": [],
   "source": [
    "tf.random.set_seed(42)\n",
    "seedstream = tfp.util.SeedStream(seed=43, salt='siwae')\n",
    "\n",
    "dataset = 'mnist'\n",
    "batch_size = 32\n",
    "num_epochs = 25\n",
    "\n",
    "data, info = tfds.load(dataset, with_info=True)\n",
    "preprocess_fn = Preprocess(info)\n",
    "\n",
    "train_data, eval_data = data['train'], data['test']\n",
    "\n",
    "train_data = tfn.util.tune_dataset(train_data,\n",
    "                                   batch_size=batch_size,\n",
    "                                   shuffle_size=1071,\n",
    "                                   preprocess_fn=preprocess_fn,\n",
    "                                   repeat_count=num_epochs)\n",
    "eval_data = tfn.util.tune_dataset(eval_data,\n",
    "                                  batch_size=500,\n",
    "                                  preprocess_fn=preprocess_fn,\n",
    "                                  repeat_count=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8HzYeAqzkgp2"
   },
   "source": [
    "# Visualizing the data\n",
    "\n",
    "Just to demonstrate that this data is _hard_ to classify, lets see what it looks like by showing a batch of data from the training set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "height": 328
    },
    "colab_type": "code",
    "id": "BoxRJqubhP8U",
    "outputId": "4635fbe9-f6d2-4061-e758-907b36fb146d"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, 112)"
      ]
     },
     "execution_count": 0,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAjQAAAElCAYAAADz6XigAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBo\ndHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAAE0BJREFUeJzt3X+s1XX9B/D3uZdlalpZErkpEUxo\nFd2imvTLbmW52moYW7iCjDRFwjvZ1WSuSamRP0LUadw16cdGQdFu1D+tlVfB1i/SO2sVWhvTFaZE\nPxZkxeV+/2h84hhf5F4v9/V+XR6Pv15/nD+ee3u5vs7zfe7ntIaHh4cLAEBiHdEBAACeKQsNAJCe\nhQYASM9CAwCkZ6EBANKz0AAA6VloAID0LDQAQHoWGgAgPQsNAJDepJG8+LnPfW6ZMmXK0cqSwgkn\nnFBKKWXv3r3BSergPNo5j3bOo53zaOc82jmPQ/vTn/5Udu3a9bSvG9FCM2XKlNLX1zfqUBNBV1dX\nKaWUwcHB4CR1cB7tnEc759HOebRzHu2cx6H19vYe0etaI/lyylarNepAE8XAwEAppZTu7u7gJHVw\nHu2cRzvn0c55tHMe7ZzHoc2ZM6ds27btaV/nMzQAQHoWGgAgPQsNAJCehQYASM9CAwCkZ6EBANKz\n0AAA6Y3owXpRnvWsZzXzv/71r8AkQDYbNmxo5gULFgQmqcPQ0FAzd3Z2BiahVgd+RrL9fGhoAID0\nUjQ0v/71r5t5+vTpgUmgfvPmzWvm/v7+wCR1mD9/fnQEYBxoaACA9Cw0AEB6Ka6cpk2bFh2Bis2a\nNauZf/Ob3wQmqcM3v/nNZu7o8J4FODb4bQcApJeioYHDue6665rZB0BLWbt2bXQEILG//vWv0RFG\nRUMDAKSXoqF55JFHoiNU5bTTTmvmP/zhD4FJ6nDeeedFR6jKpZdeGh2hKp/73OeiI0Aq55xzTnSE\nUdHQAADpWWgAgPRSXDmdfPLJ0RGqMnv27GZ25VTK8PBwdAQq9olPfCI6AqTy85//PDrCqGhoAID0\nUjQ0z3ve86IjVOXiiy9u5u9+97uBSeqwffv26AhV8e307datW9fMixcvDkwCHE0aGgAgPQsNAJBe\niisn2r33ve+NjlCVq6++OjpCVb797W8387nnnhuYpA6LFi1qZldOMHFpaACA9DQ0pNff3x8doSrv\nfOc7oyNUpdVqRUegckNDQ6WUUjo7O4OT1GHJkiWllFI+//nPBycZGQ0NAJCehQYASC/FldP+/fuj\nI1RFhc7hPPTQQ9ERquJJ0jydxx57LDpCVe68885SiisnAIBxl6KhmTVrVnSEqnjHyeG86lWvio4A\nqbzxjW+MjlCVrLciGhoAIL0UDc1vf/vb6AhV+cEPfhAdgYr985//jI4AqezYsSM6AmNAQwMApGeh\nAQDSS3Hl9PznP7+Z//znPwcmqcPu3bujIwAwQf3kJz+JjjAqGhoAIL0UDc0NN9zQzB/72McCk9Th\n4x//eHSEqhz4HpZSfBcLwDP1hje8ITrCqGhoAID0LDQAQHoprpwuvPDCZnblVMquXbuiI1TFd1u1\nu+OOO5p56dKlgUkAxo+GBgBIL0VDA4fz/ve/PzpCVZYsWdLMGhrgWKGhAQDSs9AAAOmluHJas2ZN\ndAQq1t/fHx2hKqeffnp0BCCx+fPnl1JK2bRpU3CSkdHQAADppWholi9fHh0B0vj9738fHQFIbOPG\njaWUfE9e19AAAOmlaGgARmvLli3REYBxoKEBANKz0AAA6blyAia0t73tbdERqvLwww9HR6ByHR05\nu46cqQEADqKhATiGzJo1KzoClWu1WtERRkVDAwCkZ6EBANKz0AAA6VloAID0LDQAQHoWGgAgPQsN\nAJCehQYASM9CAwCk1xoeHh4+0hfPnDmz9PX1Hc081evq6iqllDI4OBicpA7Oo53zaOc82jmPds6j\nnfM4tN7e3rJt27anfZ2GBgBIb0QNTdbvdxhLAwMDpZRSuru7g5PUwXm0cx7tnEc759HOebRzHoc2\nZ84cDQ0AcGyw0AAA6U2KDgCMrd27dzfzKaecEpgEYPxoaACA9DQ0MMF0dnZGR6Bit99+ezMvW7Ys\nMAm1GhoaKqXk+12ioQEA0rPQAADpuXKCCeZIntfAseuiiy5qZldOTCQaGgAgPQ0NTDD33XdfdISq\nfPazn23mq666KjBJHc4999zoCNV56KGHSimlnHnmmcFJ6nDbbbdFRxgVDQ0AkJ6FBgBIz5VTQqee\nemozP/HEE4FJ6nD99dc389VXXx2YpA779u2LjlCVCy64oJldOZVyzz33REeozvTp06MjVOXyyy+P\njjAqGhoAIL0UDc3pp5/ezI8++mhgkjps2bKlmV/2spcFJqnD4sWLm1lDU8rWrVujI1Rl8uTJ0RGA\ncaChAQDSs9AAAOmluHI6+GmWV155ZWCSOnhWQrsXvehF0RGq4kOfwLFIQwMApJeioZk7d250hKp0\ndNhD+f+dffbZzXzvvfcGJgEYP/7PCACkl6KhOeuss6IjVGX//v3REarSarWiI1Tl7rvvbubOzs7A\nJADjR0MDAKRnoQEA0ktx5eRDsBzOqlWroiNUZdGiRdERqrJ27droCFRueHg4OgJjwKYAAKSXoqHx\noc923//+96MjVOW6666LjlCV9evXR0eoinffPJ1vfOMb0RGq0t3dXUopZWBgIDjJyGhoAID0LDQA\nQHoprpyuueaa6AhVede73hUdoSp79+6NjkDFlixZ0sxLly4NTEKtbrrppugIVTnwsYZsz7HS0AAA\n6aVoaFavXh0dAYAJ6v7774+OwBjQ0AAA6VloAID0Ulw57dmzJzoCABwT/v73v0dHGBUNDQCQXoqG\nBmC0fJcTjMzXv/716AijoqEBANLT0JDeFVdc0cwekMVTXXrppdERqnLiiSc2s88ncigXXXRRdIRR\n0dAAAOlZaACA9FJcOc2ePbuZH3zwwcAk1OiGG25oZldOcHgPP/xwM5922mmBSWBsaWgAgPRSNDQP\nPPBAM2f79k+Ovquuuio6QlXOOOOMZn7kkUcCk9RhxYoVzbxq1arAJHXQyjBRaWgAgPQsNABAeimu\nnFqtVnQEKnbjjTdGR6jKjh07mrmjw3uW66+/vpldOXEoL37xi0sppezcuTM4SR1Wr15dSill+fLl\nwUlGxm87ACC9FA2Nd5ntDrybKMU7Cv7XwoULoyNU5fzzz4+OQOW2bt1aSillxowZwUnq0NPTU0rR\n0AAAjDsLDQCQXoorJ9r98Ic/bOaXvvSlgUmo0fr166MjVGXjxo3REajctGnToiNU5Xe/+110hFHR\n0AAA6WloEpo6dWp0BAAmqDPPPDM6wqhoaACA9DQ0CW3ZsiU6AsCE4dEgE4P/igBAehYaACA9V04J\ndXd3R0cAmDB8X+DEoKEBANKz0AAA6VloAID0LDQAQHoWGgAgPQsNAJBea3h4ePhIXzxz5szS19d3\nNPNUr6urq5RSyuDgYHCSOjiPds6jnfNo5zzaOY92zuPQent7y7Zt2572dRoaACC9ETU0Hj5UysDA\nQCnFw+0OcB7tnEc759HOebRzHu2cx6HNmTNHQwMAHBssNABAer7LifQmTfrvj/G+ffsCk9Rh/fr1\nzfzBD34wMAnA+NHQAADpaWhI78Mf/nAz33XXXYFJ6rBgwYJm1tDwVGeddVYz//jHPw5MAmNLQwMA\npGehAQDSc+VEek8++WR0BEjjwLNOSinl+OOPD0xSjwcffLCUUsrs2bODk/BMaGgAgPQ0NAlNnjy5\nmR9//PHAJHV4/etf38wH/8ky8L86OryPfaqXv/zl0REYA36yAYD0LDQAQHqunBJaunRpM19zzTWB\nSerwkY98pJl7enoCk9TBlUK7yy+/vJlvueWWwCR1OOOMM6IjVOfaa6+NjsAY8JsPAEgvRUNz8803\nN3Nvb29gkjrMnz+/mTU0pTznOc+JjlCVe++9NzpCVQ7+/aGhKeWPf/xjdITqrFy5MjoCY0BDAwCk\nZ6EBANJLceW0fPnyZnblVMqsWbOiI1CxX/ziF9ERqvKXv/wlOgIwDjQ0AEB6KRoaOJwHHnggOgIV\n+853vhMdARgHGhoAID0NDektXrw4OkJVHnvssegIVZk7d250BEhl0qT/rAb79u0LTjIyGhoAID0L\nDQCQXoorpz179kRHoGJ33313M7/gBS8ITFKHdevWRUeoyowZM6IjQCo33nhjKaX9kSkZaGgAgPRS\nNDSve93roiNQsZ/97GfREaqyc+fO6AhV+cIXvhAdgcp96lOfKqX4brwDenp6SikaGgCAcWehAQDS\nS3HltHv37ugIVWm1WtERqvKe97wnOkJV5s+f38ybNm0KTFKHSy65JDpCVS677LJmvu222wKT1OOT\nn/xkKcWV0wFf/epXoyOMioYGAEgvRUPz+OOPR0eoyvDwcHSEqgwNDUVHqMrGjRububOzMzAJNVqz\nZk0za2j+w+/UdgsXLoyOMCoaGgAgPQsNAJBeiisn4Mh96Utfio5AxT760Y9GR4CjQkMDAKSnoYEJ\nZunSpdERqNgXv/jF6AjVufbaa6MjMAY0NABAehqahDZv3hwdgYo9+eST0REglZUrV0ZHYAxoaACA\n9Cw0AEB6rpwS+sc//hEdAUhq8uTJzewp7EwkGhoAID0NDemdf/75zfy1r30tMEkdli1b1sy33357\nYBJq9KMf/aiZp0+fHpiEWt10002llFKuuOKK4CQjo6EBANKz0AAA6aW4cnr2s5/dzJ6xUcrMmTOj\nI1Rl3bp1zezKqZQ1a9Y0sysnnmratGnREajc8uXLSymunAAAxl2Khubgb4e94447ApPU4dWvfnV0\nhKq8+93vjo5QlTvvvDM6AhWbOnVqdAQq19GRs+vImRoA4CAWGgAgvRRXTjfffHMzu3IqZWBgIDpC\nVZxHu4OfQ0Mp27dvb2YfqC/l0UcfjY5QnQ0bNpRSSlmwYEFwkjq0Wq3oCKOioQEA0kvR0Bx33HHR\nEaryjne8IzoCpDFjxozoCFTOHxZMDBoaACC9FA1NX19fdAQgqax/gsr4Ofnkk6MjMAb8SwcA0rPQ\nAADppbhyWrJkSXQEIKmsf4IKjIyGBgBIz0IDAKRnoQEA0rPQAADpWWgAgPQsNABAeq3h4eHhI33x\nzJkzj/mn9nZ1dZVSShkcHAxOUgfn0c55tHMe7ZxHO+fRznkcWm9vb9m2bdvTvk5DAwCkN6KGxgOq\nShkYGCillNLd3R2cpA7Oo53zaOc82jmPds6jnfM4tDlz5mhoAIBjg4UGAEjPQgMApGehAQDSS/Ft\n2/v372/mjg47GHDkNm/e3Mzve9/7ApMAR5PtAABIz0IDAKSX4spp69at0RGqMjQ01MydnZ2BSajR\npEn//We9b9++wCR1OOWUU6IjAONAQwMApJeioTn77LOjI0Aa69ata+ZFixYFJqnDm9/85ugIwDjQ\n0AAA6VloAID0Ulw50e4DH/hAdAQq9qEPfaiZXTnxVPfcc08zv/Wtbw3LUZMTTzyxlFLKnj17gpPw\nTGhoAID0NDQJbdq0KToCFdu5c2d0BCrmjyz+19/+9rdSisdgZKehAQDSs9AAAOm5coIJpq+vLzoC\nFTv4y37hUHp6ekoppdx6663BSUZGQwMApKehgQnm+OOPj45AxTo6vI/l8FavXl1K0dAAAIw7DQ1M\nMJdddlkzr1ixIjAJNfIZGiYqDQ0AkJ6FBgBIL8WVk+8egSPnQ8Eczre+9a3oCNXp7e2NjsAY0NAA\nAOmlaGje8pa3REegYvPmzWvm/v7+wCRQv0suuSQ6QnVuueWW6AhV2b59e3SEUdHQAADpWWgAgPRS\nXDkdd9xx0RGo2Fe+8pVmPumkkwKTQP2eeOKJ6AhU7rWvfW10hFHR0AAA6aVoaP79739HR6BiJ5xw\nQnSEqqxfvz46ApDY3r17oyOMioYGAEjPQgMApJfiygkOZ8OGDdERqrJw4cLoCADjTkMDAKSnoSG9\nL3/5y9ERqvKKV7yimX/5y18GJgEYPxoaACA9DQ3pfe9734uOUJWDv5fmnHPOCUwCOXR2dpZSShka\nGgpOUocDj8LI9ufbGhoAID0LDQCQnisnmGDe/va3R0eAVC644IJSSil33XVXbJBKLFq0qJRSytq1\na4OTjIyGBgBIT0OT0Jve9KZmvu+++wKTUKNPf/rT0RGq0tfX18wXX3xxYBJqdeWVV5ZSNDQHHDgP\nDQ0AwDiz0AAA6aW4curq6mrmwcHBwCR12LRpUzNPmTIlMAk1WrlyZXSEqlx44YXN7MqplBUrVjTz\nqlWrApPUY8aMGdERqjJ16tToCKOioQEA0kvR0Nx///3N3NFhBzv11FOjI1Tl1ltvbeaenp7AJFC/\n8847r5k1NEwktgMAID0LDQCQXoorp82bN0dHoGLLli1rZldOPFWr1YqOUJXXvOY10RGo3K9+9avo\nCKOioQEA0kvR0MybNy86QlXmzp0bHaEqPijO4fj5aLd9+/boCFTula98ZXSEUfEvHQBIL0VDQ7uf\n/vSn0RGApD7zmc9ER4CjQkMDAKRnoQEA0nPlBHAM6e/vj45Qnb1790ZHYAxoaACA9DQ0AMeQPXv2\nREeozkknnRQdgTGgoQEA0rPQAADpWWgAgPRaw8PDw0f64he+8IXlJS95yVGMAwDwXzt27Ci7du16\n2teNaKEBAKiRKycAID0LDQCQnoUGAEjPQgMApGehAQDSs9AAAOlZaACA9Cw0AEB6FhoAIL3/AwOR\n3beP1RXBAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {
      "tags": []
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def reshape_image_batch(tensor, cut=None, rows=None, axis=0):\n",
    "  \"\"\"Given an array of shape [n, x, y, ...] reshape it to create an image field.\n",
    "\n",
    "  Arguments:\n",
    "    tensor: The array to reshape.\n",
    "    cut: Optional cut on the number of images to view. Will default to whole\n",
    "      array.\n",
    "    rows: Number of rows to use.  Will default to the integer less than the\n",
    "      sqrt.\n",
    "    axis: Axis to interpretate at the batch dimension.  By default the image\n",
    "      dimensions immediately follow.\n",
    "\n",
    "  Returns:\n",
    "    reshaped_array: An array of shape [rows * x, cut / rows * y, ...]\n",
    "  \"\"\"\n",
    "  original_shape = tensor.get_shape().as_list()\n",
    "  assert len(original_shape) >= 2, 'array must be at least 3 Dimensional.'\n",
    "\n",
    "  if cut is None:\n",
    "    cut = original_shape[axis]\n",
    "  if rows is None:\n",
    "    rows = int(np.sqrt(cut))\n",
    "\n",
    "  cols = cut // rows\n",
    "  cut = cols * rows\n",
    "\n",
    "  leading = original_shape[:axis]\n",
    "  x_width = original_shape[axis + 1]\n",
    "  y_width = original_shape[axis + 2]\n",
    "  remaining = original_shape[axis + 3:]\n",
    "\n",
    "  tensor = tensor[..., :cut, :, :, :]\n",
    "  tensor = tf.reshape(tensor, leading + [rows, cols, x_width, y_width] + remaining)\n",
    "  new_shape = [i for i in range(len(tensor.get_shape().as_list()))]\n",
    "  new_shape[axis + 1] = axis + 2\n",
    "  new_shape[axis + 2] = axis + 1\n",
    "  tensor = tf.transpose(tensor, new_shape)\n",
    "  tensor = tf.reshape(tensor, leading + [rows * x_width, cols * y_width] + remaining)\n",
    "\n",
    "  return tensor\n",
    "\n",
    "x, y = next(iter(train_data))\n",
    "xim = tfb.Pad(paddings=((13, 14),), constant_values=-1.)(x[..., tf.newaxis])\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.imshow(reshape_image_batch(xim[..., tf.newaxis], rows=4)[:,:,0],cmap='gray')\n",
    "for i in range(8):\n",
    "  plt.axhline(i*28-0.5, color='w')\n",
    "  plt.axvline(i*28-0.5, color='w')\n",
    "plt.xticks([])\n",
    "plt.yticks([])\n",
    "plt.xlim(0, 8*28)\n",
    "plt.ylim(0, 4*28)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZlkBTS5wKEyG"
   },
   "outputs": [],
   "source": [
    "class Encoder(tf.Module):\n",
    "  def __init__(self, input_size=28,\n",
    "               num_hidden=128, \n",
    "               encoded_size=2, \n",
    "               nmix=2, \n",
    "               name=None):\n",
    "\n",
    "    super(Encoder, self).__init__(name=name or 'encoder')\n",
    "\n",
    "    self.input_size = input_size\n",
    "    self.nmix=nmix\n",
    "    self.num_hidden=num_hidden\n",
    "    self.encoded_size=encoded_size\n",
    "\n",
    "    # Convenience function\n",
    "    affine = functools.partial(tfn.Affine,\n",
    "                               init_kernel_fn=tf.initializers.he_normal(),\n",
    "                               init_bias_fn = tf.zeros_initializer())\n",
    "\n",
    "    self._dnn = tfn.Sequential([\n",
    "        affine(self.input_size, self.num_hidden), tf.nn.elu,\n",
    "        affine(self.num_hidden, self.num_hidden), tf.nn.elu,\n",
    "        affine(self.num_hidden, self.num_hidden), tf.nn.elu,\n",
    "        affine(self.num_hidden, self.num_hidden),\n",
    "        affine(self.num_hidden, self.num_out)])\n",
    "\n",
    "  def __call__(self, inputs):\n",
    "    outputs = self._dnn(inputs)\n",
    "\n",
    "    loc_bijector = tfb.Reshape([self.nmix, self.encoded_size])\n",
    "    scale_bijector = tfb.FillScaleTriL()(\n",
    "        tfb.Reshape([self.nmix,self.encoded_size*(self.encoded_size+1)//2]))\n",
    "\n",
    "    return tfd.MixtureSameFamily(\n",
    "        tfd.Categorical(outputs[..., :self.nmix]),\n",
    "        tfd.MultivariateNormalTriL(\n",
    "            loc=loc_bijector(\n",
    "                outputs[..., self.nmix:self.nmix*(1+self.encoded_size)]),\n",
    "            scale_tril=scale_bijector(\n",
    "                outputs[..., self.nmix*(1+self.encoded_size):])))\n",
    "\n",
    "  @property\n",
    "  def num_out(self):\n",
    "    return self.nmix*(1 + self.encoded_size +\n",
    "                      self.encoded_size * (self.encoded_size + 1) // 2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3SdEstSXPQR7"
   },
   "outputs": [],
   "source": [
    "class Decoder(tf.Module):\n",
    "  def __init__(self, encoded_size=2, num_classes=10, name=None):\n",
    "    super(Decoder, self).__init__(name=name or 'decoder')\n",
    "\n",
    "    self.encoded_size = 2\n",
    "    self.num_classes = 10\n",
    "\n",
    "    # Convenience function\n",
    "    affine = functools.partial(tfn.Affine, \n",
    "                               init_kernel_fn=tf.initializers.he_normal(),\n",
    "                               init_bias_fn = tf.zeros_initializer())\n",
    "\n",
    "    # DNN is just an affine transformation for the decoder\n",
    "    self._dnn = affine(self.encoded_size, self.num_classes)\n",
    "\n",
    "  def __call__(self, inputs):\n",
    "    return tfd.Categorical(logits=self._dnn(inputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fgVCUq-nTcIW"
   },
   "outputs": [],
   "source": [
    "class Vae(tf.Module):\n",
    "  def __init__(self, encoder, decoder, prior, nz, beta=1.0, seed=None):\n",
    "    self.encoder = encoder\n",
    "    self.decoder = decoder\n",
    "    self.prior = prior\n",
    "    self.nz = nz\n",
    "    self.beta = beta\n",
    "    self.seed=seed\n",
    "\n",
    "  def __call__(self, x, y):\n",
    "\n",
    "    # Surrogate posterior\n",
    "    pzgx = self.encoder(x)\n",
    "\n",
    "    # Posterior component weights\n",
    "    mix_logits = pzgx.mixture_distribution.logits_parameter()  # [n, k]\n",
    "\n",
    "    # Posterior samples\n",
    "    z = pzgx.components_distribution.sample(self.nz,\n",
    "                                            seed=self.seed)  # [t, n, k, d]\n",
    "\n",
    "    # Likelihood distribution\n",
    "    pygz = self.decoder(z)\n",
    "\n",
    "    # log p(y|z)\n",
    "    lpygz = pygz.log_prob(y[..., tf.newaxis])  # [t, n, k]\n",
    "\n",
    "    # log q(z|x), which requires two annoying transposes\n",
    "    transpose_second_and_third_dims = lambda x: tf.transpose(\n",
    "        x, tf.concat([[0, 2, 1], tf.range(3, tf.rank(x))], axis=0))\n",
    "\n",
    "    lqzgx = transpose_second_and_third_dims(\n",
    "        pzgx.log_prob(transpose_second_and_third_dims(z)))  # [t, n, k]\n",
    "\n",
    "    # log r(z)\n",
    "    lrz = self.prior.log_prob(z)  # [t, n, k]\n",
    "\n",
    "    return lpygz, lqzgx, lrz, mix_logits\n",
    "\n",
    "class Selbo(Vae):\n",
    "  def __call__(self, x, y):\n",
    "\n",
    "    # Get everything you need from the base class\n",
    "    lpygz, lqzgx, lrz, mix_logits = super(Selbo, self).__call__(x, y)\n",
    "    alpha = tf.math.softmax(mix_logits, axis=-1)\n",
    "\n",
    "    # for housekeeping, lets place this in terms of nll and kld\n",
    "    K = tf.cast(tf.shape(alpha)[-1], tf.float32)\n",
    "    nll = -tf.reduce_mean(lpygz * alpha, axis=[0, 1, 2]) * K\n",
    "    kl = tf.reduce_mean((lqzgx - lrz) * alpha, axis=[0, 1, 2]) * K\n",
    "\n",
    "    return nll + beta * kl, (nll, kl)\n",
    "\n",
    "class Siwae(Vae):\n",
    "  def __call__(self, x, y):\n",
    "    # Get everything you need from the base class\n",
    "    lpygz, lqzgx, lrz, mix_logits = super(Siwae, self).__call__(x, y)\n",
    "\n",
    "    # Convert logits into log-probs\n",
    "    log_alpha = tf.math.log_softmax(mix_logits, axis=-1)\n",
    "\n",
    "    # For comparison to SELBO, lets gather the nll and kld\n",
    "    alpha = tf.math.softmax(mix_logits, axis=-1)\n",
    "    K = tf.cast(tf.shape(alpha)[-1], tf.float32)\n",
    "    nll = -tf.reduce_mean(lpygz * alpha, axis=[0, 1, 2]) * K\n",
    "    kl = tf.reduce_mean((lqzgx - lrz) * alpha, axis=[0, 1, 2]) * K\n",
    "\n",
    "    # Importance weights:  shape = [t, n, k]\n",
    "    wts = lpygz + self.beta * (lrz - lqzgx) + log_alpha\n",
    "\n",
    "    # LME over samples, LSE over modes:  shape = [n]\n",
    "    T = tf.cast(tf.shape(wts)[0], tf.float32)\n",
    "    siwae = tf.math.reduce_logsumexp(wts, axis=[0, -1]) - tf.math.log(T)\n",
    "\n",
    "    # SIWAE is average over data:  shape = []\n",
    "    return -tf.math.reduce_mean(siwae, axis=0), (nll, kl)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Gm3hq3JnlFjn"
   },
   "source": [
    "# Define the model\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Kuj8GG8_WFt_"
   },
   "outputs": [],
   "source": [
    "encoded_size = 2\n",
    "num_hidden = 128\n",
    "nmix = 10\n",
    "nz = 1\n",
    "beta = 0.05\n",
    "lr = 0.001\n",
    "\n",
    "# ---------------------------------------------------------------\n",
    "\n",
    "loss_fn = Selbo(\n",
    "    encoder=Encoder(input_size=info.features['image'].shape[0],\n",
    "                    encoded_size=encoded_size,\n",
    "                    num_hidden=num_hidden,\n",
    "                    nmix=nmix),\n",
    "    decoder=Decoder(encoded_size=encoded_size,\n",
    "                    num_classes=info.features['label'].num_classes),\n",
    "    prior=tfd.MultivariateNormalDiag(loc=tf.zeros(encoded_size),\n",
    "                                     scale_diag=tf.ones(encoded_size)),\n",
    "    nz=nz,\n",
    "    beta=beta,\n",
    "    seed=seedstream)\n",
    "\n",
    "opt = tf.optimizers.Adam(lr)\n",
    "train_op = tfn.util.make_fit_op(\n",
    "    loss_fn=loss_fn, optimizer=opt,\n",
    "    trainable_variables=loss_fn.trainable_variables,\n",
    "    xla_compile=False)  # xla doesn't like this... not sure why"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ww92nDQZmNz8"
   },
   "source": [
    "# Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "height": 197
    },
    "colab_type": "code",
    "id": "kNXQyI5PWdie",
    "outputId": "cf6dd95e-6484-4e4c-eac0-c223b88b50ed"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iteration 0,  loss:  7.0789,   nll:  6.5202,   kl:  11.1747\n",
      "iteration 5000,  loss:  1.5663,   nll:  1.4026,   kl:  3.2740\n",
      "iteration 10000,  loss:  1.1439,   nll:  0.9116,   kl:  4.6454\n",
      "iteration 15000,  loss:  0.9195,   nll:  0.7426,   kl:  3.5390\n",
      "iteration 20000,  loss:  1.3499,   nll:  1.1602,   kl:  3.7941\n",
      "iteration 25000,  loss:  0.8219,   nll:  0.6056,   kl:  4.3245\n",
      "iteration 30000,  loss:  0.8417,   nll:  0.6280,   kl:  4.2750\n",
      "iteration 35000,  loss:  1.1429,   nll:  0.9326,   kl:  4.2060\n",
      "iteration 40000,  loss:  1.2736,   nll:  1.0895,   kl:  3.6814\n",
      "iteration 45000,  loss:  0.9169,   nll:  0.7086,   kl:  4.1654\n"
     ]
    }
   ],
   "source": [
    "for i, (x, y) in enumerate(iter(train_data)):\n",
    "  this_loss, (this_nll, this_kl) = train_op(x, y)\n",
    "  if i %5000 == 0:\n",
    "    print(\"iteration %i,  loss:  %.4f,   nll:  %.4f,   kl:  %.4f\"%(\n",
    "        i, this_loss.numpy(), this_nll.numpy(), this_kl.numpy()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GJNArpbllBo5"
   },
   "source": [
    "# Evaluate the accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7z0qZpTHdufZ"
   },
   "outputs": [],
   "source": [
    "def evaluate_accuracy(encoder, decoder, dataset, nz=100, seed=None):\n",
    "  \"\"\"Evaluate the accuracy of a model on a dataset.\"\"\"\n",
    "  it = iter(dataset)\n",
    "  hits = 0\n",
    "  trials = 0\n",
    "\n",
    "  for x, y in it:\n",
    "    pzgx = encoder(x)\n",
    "    z = pzgx.sample(nz, seed=seed)\n",
    "    pygz = decoder(z)\n",
    "    # average py over samples / components and get argmax as prediction\n",
    "    avg_log_py = tfp.math.reduce_logmeanexp(\n",
    "        tf.math.log_softmax(pygz.logits_parameter(), axis=-1), axis=0)\n",
    "    yhat = tf.cast(tf.argmax(avg_log_py, axis=-1), tf.int32)\n",
    "    hits += tf.reduce_sum(tf.cast(tf.equal(yhat, y), tf.int32))\n",
    "    trials += tf.size(y)\n",
    "  acc = tf.cast(hits, tf.float32)/tf.cast(trials, tf.float32)\n",
    "  return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "height": 35
    },
    "colab_type": "code",
    "id": "CcDC1VGGfwMn",
    "outputId": "d83230f6-4806-4f72-da46-08d1f23b6a25"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy:   0.7262\n"
     ]
    }
   ],
   "source": [
    "# Evaluate predictive accuracy with 10k samples\n",
    "print(\"accuracy:   %.4f\"%(\n",
    "    evaluate_accuracy(encoder=loss_fn.encoder,\n",
    "                      decoder=loss_fn.decoder,\n",
    "                      dataset=eval_data,\n",
    "                      nz=10000,\n",
    "                      seed=seedstream).numpy()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QSUZbep9OW6R"
   },
   "source": [
    "# Comparing performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "height": 306
    },
    "colab_type": "code",
    "id": "-CdO0hStOVe0",
    "outputId": "41ef83b9-a39f-447e-9b0f-9d671174a95b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7fb3fbcd0a20>"
      ]
     },
     "execution_count": 0,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAZEAAAEPCAYAAACDTflkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBo\ndHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAAIABJREFUeJzt3XlcVPX6wPEPi7vihqiIgoYigooC\nmpk7oje3SivNckvJzLXrLZeukaX500wt783oulUIZanIEkWuqaiRS5ob7kBmhIoisZ/fHydGkRkZ\nYIYzwPN+veaFc+bMnIdR55nv9nytFEVREEIIIUrAWusAhBBClF+SRIQQQpSYJBEhhBAlJklECCFE\niUkSEUIIUWKSRIQQQpSYJBEhhBAlJklECCFEiUkSEUIIUWK2Wgdgbvb29ri4uGgdhhBClBuXL1/m\nzz//NOrcCp9EXFxciIuL0zoMIYQoN3x8fIw+V7qzhBBClJgkESGEECUmSUQIIUSJVfgxEX2ys7NJ\nTEwkIyND61DE36pXr46TkxNVqlTROhQhRDFUyiSSmJhInTp1cHFxwcrKSutwKj1FUUhJSSExMZGW\nLVtqHY4QohgqZXdWRkYGDRs2lARiIaysrGjYsGG5aBkGB4OLC1hbqz+Dg7WOSAhtVcqWCCAJxMKU\nh7+P4GAICID0dPX+lSvqfYDRo7WLSwgtlXlLJDo6Gjc3N1xdXVmyZEmhx5ctW4aXlxdeXl54enpi\nY2PDjRs3OHv2rO64l5cXdnZ2rFy5sqzDF5XY/Pn3Eki+9HSYNQt274ajR+HiRUhJgZwcTUIUosxZ\nleUe67m5ubRp04aYmBicnJzw9fUlJCSEdu3a6T0/PDycFStWsHPnzkKv06xZMw4dOoSzs/NDr+nj\n41NoseHp06dxd3cv3S9TSosWLWLTpk3Y2NhgbW3NJ598QteuXQ2eP27cOAYPHsyIESN0Cyjt7e1L\ndO1t27bxyy+/sGDBgpKGb9DmzZsJDAzk9OnTHD58WLdo6cSJEyxfvpwNGzYYfK4l/L08jLU1FOd/\nS82aULdu6W4yz0BoQd/npiFl2p11+PBhXF1dadWqFQAjR44kLCzMYBIJCQlh1KhRhY7v2LGDRx55\npMgEYqliY2OJiIjgyJEjVKtWjT///JOsrKwyu/7SpUvZvn27WV7b09OTLVu28PLLLxc43r59exIT\nE7l69SotWrQwy7XN6fZtqF4d/vqr8GNNmsCmTZCaWvTt6tV7f36wVaNPjRqlT0RVq5r+/RAiX5km\nkaSkJJo3b6677+TkxKFDh/Sem56eTnR0NKtXry70WGhoqN7kUl5cu3YNe3t7qlWrBlCgRfHzzz/z\n2muvkZaWhr29PRs2bKBp06aFXmPZsmXs2rULgE2bNuHq6sqVK1eYMGECycnJNGrUiPXr1xf6wD53\n7hzVqlXTXXPcuHHY2dkRFxfH77//ztKlSxkxYgQLFizQJZrk5GT8/f1Zv359kb/bw1oSQ4YMITQ0\nlNdff73I17Ekp07BU09BRobaMsjOvvdYzZrw/vvQp0/xXzc7W01OxiSfW7fu/TkhQRKRsBxlmkT0\n9ZwZGlANDw+ne/fuNGjQoMDxrKwstm/fznvvvWfwOkFBQQQFBQHqB+BDzZwJx44VEXkxeXnBQ8Zr\n/P39WbhwIW3atMHPz4/nnnuOXr16kZ2dzbRp0wgLC6NRo0Z8+eWXzJ8/n3Xr1hV6DTs7Ow4fPsxn\nn33GzJkziYiIYOrUqYwZM4axY8eybt06pk+fzrZt2wo8b//+/XTu3LnAsWvXrrFv3z7OnDnD0KFD\nGTFiBAsXLmThwoWkpqbSo0cPpk6dCkCPHj24c+dOoXjef/99/Pz8Hvq2+Pj4sGTJknKVRL75BsaN\nU5PF7t3qB/j8+WqLokULWLSo5IPqVapAw4bqraSKk4juvyUm3vvz3btFX6d69dInor+/M4kKpkyT\niJOTEwkJCbr7iYmJODo66j3XUGvj22+/pXPnzjRu3NjgdQICAgj4e9pMcQqJlZXatWvz888/8+OP\nP7Jr1y6ee+45lixZgo+PDydPnqR///6AOvajrxUC6N6bUaNGMWvWLEDtJtuyZQsAL774ot4P62vX\nrtGoUaMCx5588kmsra1p164d169f1x1XFIXRo0cza9YsvL29Afjxxx9L/Hs7ODjw22+/lfj5ZSkn\nR00WS5fCo4/C119Ds2bqY5Y0E8sUiSgnp2SJ6Lff7v05La3o61SrVvpEVL16yX9PYR5lmkR8fX2J\nj4/n0qVLNGvWjNDQUDZt2lTovNTUVPbs2cMXX3xR6DFD4yQlptEMLxsbG3r37k3v3r1p3749Gzdu\nxNvbGw8PD2JjY4t8/v0tOEOtOX3Ha9SoQWpqaoFj1e77inh/azEwMBAnJyfGjx+vO1aalkhGRgY1\natR46DmW4M8/YdQo+OEHmDxZ/SdSkb9F29pCgwbqraRyc/Unovu74fTdrl0rXiKqWtU0iagczCgv\nN8o0idja2rJ69WoGDBhAbm4uEyZMwMPDgzVr1gAwefJkALZu3Yq/vz+1atUq8Pz09HRiYmL45JNP\nyjJskzt79izW1ta0bt0agGPHjuHs7IybmxvJycnExsbSrVs3srOzOXfuHB4eHoVe48svv2TOnDl8\n+eWXdOvWDYDHHnuM0NBQXnzxRYKDg3n88ccLPc/d3V1vcn5QREQEMTEx7N69u8Dx0rREzp07h6en\nZ4mfXxZ+/hmefhquX4e1a2HCBK0jKh9sbKB+ffVWUoYSUVG369fv/VnP95tCqlQpmFTq1ZNEVBpl\nvtjwiSee4IknnihwLD955Bs3bhzjxo0r9NyaNWuSkpJizvDKRFpaGtOmTePWrVvY2tri6upKUFAQ\nVatW5euvv2b69OmkpqaSk5PDzJkz9SaRzMxMunbtSl5eHiEhIQB8+OGHTJgwgWXLlukG1h/Us2dP\n/vnPf6IoykMX+C1fvpzffvuNLl26ADB06FAWLlxY5O+2detWpk2bRnJyMoMGDcLLy4vvvvsOgF27\ndjFo0CCj3iMtbNigtjwaN4Z9+8ACe0IrNFMlojt3ip+I4uPv/fn27aKv82AiKsmtRg3zJKLgYNON\n2xnDqHUiqamp1K1b13xRmJGlrhPR0owZMxgyZEiR3U+mlJmZSa9evdi3bx+2tvq/u2j195KVpc6v\n+Phj6NcPQkLggWEjUYnk5amJqKiuuIfd7twpek2RrW3pE1HNmgUT0YNVFUA9JyioeInE5OtEHB0d\nGTlyJJMnT8bX19f4SIRFmjdvnsGp1eZy9epVlixZYjCBaCUpCUaMgIMH4fXX1W9tFhaiKGPW1vc+\npEsqPxEVN/lcvFiwRWRMIrKzuxfv6dOQmVnwnPR0tWVirtaIUf9d/vWvf7Fu3To2bNhAhw4deOWV\nV3j++eepXbu2eaISZtW4cWOGDh1aptds3bq1bgzIUuzdC88+qw7obt6sJhMhTMFUiSgtrXhJyNBq\nhatXSx5HkRQj5ebmKmFhYcoTTzyh2NjYKHXq1FEmT56sHD161NiX0IS3t3ehY6dOndIgElGUsvp7\nyctTlFWrFMXWVlHatFGUX38tk8sKYXbOzoqitl8K3pydi/c6+j43DTG6AKO1tTVDhw4lMjKSCxcu\nMGPGDLZv3463tzddu3Zlw4YNZD7YjhLCwqSnw4svwowZMGgQHD4MBqruCFHuLFqkjoHcr2ZN9bi5\nlKiKr52dHQ0aNKB27dooikJqaiovvfQSrq6u7Nu3z9QxCmESFy/CY4+pda7efRe2bCldd4MQlmb0\naHUQ3dlZHXB3di7+oHpxFSuJ7N+/nzFjxtCsWTPeeust+vbty/Hjxzlz5gynT5+mVatWhQrvCWEJ\nvv0WvL3VvuGoKHWg0bpSbskmKrrRo+HyZXVM5fJl81dYMOq/0UcffYSnpyc9e/bkyJEjLFu2jKSk\nJD7++GPat28PQJs2bXj77bc5c+aMWQOuKBYtWoSHhwcdOnTAy8tLN1tq4sSJnDp1qkxi2LZtm1Fr\nP0pi8+bNeHh4YG1tXWCq4IkTJ/SuATKXvDy11TFokPqtLC4OBg4ss8sLUeEZNTtr9uzZPPnkk/zn\nP/+hV69eBs9r3bq1WfaoqGgeVgr+f//7X5nFUdFLwqemwpgxsH37vWb+g/3FQojSMaolcvXqVb78\n8suHJhBA180lHk5fKfj8QpS9e/cmLi6Or776itdeew2AVatW6fZguXDhgq6cycKFC/H19cXT05OA\ngABd3asLFy4wcOBAvL296dGjh97Wob6S8NOnT+exxx6jVatWfP311wAsWLBAt5tks2bNCtTRehh3\nd3fc3Nz0PpZfEt6cfv0VfH3VrqsPP4TPP5cEIoQ5GNUSSU1N5cyZM3qTyN69e2natKnFrQEwlgaV\n4A2Wgr9fz549WbZsGaDWq2rYsCFJSUns27ePHj16ADB16lRdy+/FF18kIiKCIUOGEBAQwJo1a2jd\nujWHDh1iypQphXaHrMgl4TdvhvHjoXZt2LkT/n67hBBmYFQSmTlzJu3atdObRCIiIjh16hQREREm\nD66iMlQK/v6xgiZNmpCWlsadO3dISEjg+eefZ+/evfz44488/fTTgFqLaunSpaSnp3Pjxg08PDzo\n06cPBw4c4JlnntG9lr6p1xWxJHxODsybB8uWQbduavl2AzsNCCFMxKgkEhcXV6hIYr6ePXuyceNG\nkwZVljSqBK+3FPyDA87dunVj/fr1uLm50aNHD9atW0dsbCzLly8nIyODKVOmEBcXR/PmzQkMDCQj\nI4O8vDzq1avHsSKaVxWtJHxyMowcqbY8pkyBFStkNz4hyoJRSeTOnTtUN7AbTJUqVQp9GImHM1QK\n/kE9e/ZkwYIFLFiwgE6dOrFr1y5q1KhB3bp1uXXrFqCOp6SlpfH1118zYsQI7OzsaNmyJZs3b+aZ\nZ55BURR++eUXOnbsWOC1K1JJ+Lg4tXz7H3/A+vXqToRCiLJh1MB6q1at2LFjh97Hdu7ciYuLiylj\nqvDS0tIYO3Ys7dq1o0OHDpw6dYrAwMBC5/Xo0YOEhAR69uyJjY0NzZs31w2q16tXj0mTJtG+fXue\nfPLJAoUxg4ODWbt2LR07dsTDw4OwsLBCr92zZ0+OHj2qd8vi+91fEt7Ly8vo2Xdbt27FycmJ2NhY\nBg0axIABA3SPmbIk/Lp18Pjj6pqP/fslgQhR5oypjbJkyRKlSpUqyurVq5WMjAxFURQlIyNDWb16\ntVK1alVl8eLFxSvMUoakdpZh06dPV2JiYsr0mhkZGUrXrl2V7OzsQo8V5+8lI0NRXn5ZrQvk56co\nycmmjFKIys3ktbNmz57N0KFDmTZtGrVq1cLBwYFatWoxbdo0hg4dyhtvvGHuXCfMYN68eaTfv/FA\nGTBFSfjEROjVCz75BObMgeho+HumshCijBn1P9nGxoavv/6anTt3EhMTQ0pKCvb29vj7+9O7d28z\nhyjMpTyWhN+zRy3fnp6uzr4aPtyEwQkhiq1YXwf79u1L3759zRVLmVKK2B5WlC2liLEZRYFVq2D2\nbHB1hd27oRJvTimExaiUJeiqV69OSkpKkR9comwoikJKSorBGYB376plS2bNgiFD1PLtkkCEsAxG\nt0SCgoL4+OOPOXv2rN7Fa7m5uSYNzJycnJxITEwkOTlZ61DE36pXr46Tk1Oh4xcuwFNPwcmTsHgx\nvPGGVN8VwpIYlUQ+++wzpk2bxtixYzl+/DgTJkwgOzub7du306hRI0abu9awiVWpUoWWLVtqHYYo\nQlSU2gKxtlZLud83S1gIYSGM+k63cuVK5s6dy8cffwzAlClT2LhxIxcvXqRGjRo0bNjQrEGKyiUv\nDxYuhMGDwcVFXUwoCUQIy2RUEomPj6dnz55YW1tjbW2tK1tev3595s+fz6pVq8wapKg8bt2CYcPg\nrbfghRfUBYTSaBTCchmVRGrUqEFeXh5WVlY0adKEixcv6h6rXbu2WYrpicrn5Em1fHt0NHz0EWzc\nKOXbhbB0Ro2JtG/fnvPnz+Pn50ePHj1YvHgxLVu2xNbWlsDAQNq2bWvuOEUF9+WXMGEC2Nmp03e7\nd9c6IiGEMYxqiQQEBHDz5k0A3nnnHdLS0nj88cd59NFHOXfuHMuXLzf6gtHR0bi5ueHq6sqSJUsK\nPb5s2TLdJkienp7Y2Nhw48YNAG7dusWIESNo27Yt7u7uxMbGGn1dYZlyctS1HyNHQqdOcOSIJBAh\nypWS1FVJS0tTYmJilLCwMCW5GEWLcnJylFatWikXLlxQMjMzlQ4dOii//vqrwfO3b9+u9OnTR3d/\nzJgxyqeffqooiqJkZmYqN2/eLPKaxakBI8rW9euK0qePWv9q6lRFyczUOiIhhKKYuHZWVlYWs2bN\n4qefftIdq1WrFn5+fgwdOlS3vaoxDh8+jKurK61ataJq1aqMHDlSb4XZfCEhIYwaNQqA27dvs3fv\nXl566SUAqlatSr169Yy+trAshw+DtzfExsKGDeoYiOz/IUT5U2QSqVq1Kp988gl//fVXqS+WlJRE\n8+bNdfednJxISkrSe256ejrR0dEM/7s40sWLF2nUqBHjx4+nU6dOTJw4kbt375Y6JlH2/vc/dcta\nW1s4cADGjtU6IiFESRk1JtKpUydOnDhR6ospesqMGKpfFR4eTvfu3WnQoAEAOTk5HDlyhFdeeYWj\nR49Sq1YtvWMqoK6u9/HxwcfHR1alW5DMTAgIgEmToHdvdf1Hp05aRyWEKA2jksjy5ct5//33iYiI\nKFW9KScnJxISEnT3ExMTcTSwCXZoaKiuKyv/uU5OTnTt2hWAESNGcOTIEb3PDQgIIC4ujri4uEL7\niAttJCRAz57w6acwd666Gl3WqApR/hk1xfeZZ54hNTWVYcOGYWtri4ODQ4EWhJWVFVeuXCnydXx9\nfYmPj+fSpUs0a9aM0NBQNm3aVOi81NRU9uzZU2D71iZNmtC8eXPOnj2Lm5sbO3bsoF27dsaELzS2\ne7davj0jA7ZsUWthCSEqBqOSSL9+/UxSNt3W1pbVq1czYMAAcnNzmTBhAh4eHqxZswaAyZMnA+rW\nqv7+/tSqVavA8z/66CNGjx5NVlYWrVq1Yv369aWOSZiPosCKFfD669C6NWzdCrKkSIiKxUopTf9U\nOeDj40NcXJzWYVQ6d+/CSy+piwifekqdgWVnp3VUQghjFOdzU4pqC5M7fx4efRQ2b4b33oNvvpEE\nIkRFZXQp+KKMGTOm1MGI8i8iQi2caGOj1sDq31/riIQQ5mRUEhk3bpze4/ePk0gSqdzyy7e//bY6\nbXfLFrWMuxCiYjMqiVy6dKnQsZSUFCIiIti0aVOBWVSi8rl1S219REbCmDGwZg3UqKF1VEKIsmBU\nEnF2dtZ7rHPnziiKwgcffKB3qq6o+E6cUAfOr1yB//wHXnkFTDCRTwhRTpR6YL1Hjx5ERkaaIhZR\nzoSGqgPo6emwZw9MmSIJRIjKptRJ5ODBg9SuXdsUsYhyIjsbXnsNRo2Czp3h55/hsce0jkoIoQWj\nurMWLlxY6FhWVhYnT54kMjKSqVOnmjwwYZmuX4fnnlNbHtOmwfvvS/VdISozo5JIYGBgoWPVqlXD\n2dmZ+fPnM3fuXFPHJSzQoUMwfDikpMBnn8GLL2odkRBCa0Ylkby8PHPHISzcp5/C1Kng6KiWb5fq\nu0IIkBXroggZGWrp9oAA6NNHHf+QBCKEyGdUElm/fr3eLi1Qu7o2btxoypiEhcgv3/6//8H8+eo6\nkL+3dxFCCMDIJLJq1SoaGtj8wcHBgZUrV5o0KKG9XbvUmVdnzqjVd999Vy1lIoQQ9zMqiZw/fx4P\nDw+9j7m7u3PhwgWTBiW0oyjqjCs/P2jUCH76CZ58UuuohBCWyqiBdVtbW/7880+9j8n2sxVHWppa\nvv2rr9RZWOvXQ506WkclhLBkRrVEunTpots46kFr1qzB19fXpEGJshcfr64+//pr+L//U8u4SwIR\nQhTFqJbI/Pnz8fPzo2vXrkycOJFmzZqRlJTE//73P44cOUJMTIy54xRmFB6uFlCsUgW++07tyhJC\nCGMYlUR69erF119/zcyZM3n55Zd1x11cXPjmm2/o3bu3ueITZpSXp5ZuX7hQHUTfsgX01NoUQgiD\njEoiAMOGDWPYsGGcPXuWlJQU7O3tadOmjTljE2Z086ba+oiKgnHj4L//lfLtQojiMzqJ5HNzczNH\nHKIM/fKLWr49IQE+/hheflmq7wohSsaogfVZs2bxooFCSS+++CKzZ882aVDCfDZtUgfQMzLUIoqT\nJ0sCEUKUnFFJZPv27fj7++t9bMCAAWzbts2kQQnTy86GWbNg9Gjw8VHLl3TrpnVUQojyzqjurKSk\nJJo3b673MScnJ5KSkkwalDCt69fh2Wdh716YPl1dTFilitZRCSEqAqNaIvXr1+f8+fN6Hzt//jx1\nZEGBxTp4UJ159dNP8PnnsGqVJBAhhOkYlUT8/PxYtGgR169fL3D8+vXrLF68mP79+5slOFFyigKf\nfKIWUKxWDWJj1dlYQghhSkZ1Z73zzjv4+vrSunVrBg8erOvCioiIoFq1arz77rvmjlMUQ0YGvPoq\nrFsHAwdCcLBU3xVCmIdRScTFxYWffvqJBQsWEBMTo1sn8tRTT/H222/jLCvULMbVq2rdq7g4ePNN\nCAyU6rtCCPMxep2Ii4sLn332WakvGB0dzYwZM8jNzWXixInMmTOnwOPLli0jODgYgJycHE6fPk1y\ncjINGjTAxcWFOnXqYGNjg62tLXFxcaWOpyLZsQNGjoSsLNi2DYYN0zoiIUSFp5RCWlqasmHDBqVP\nnz5GnZ+Tk6O0atVKuXDhgpKZmal06NBB+fXXXw2ev3379gKv7ezsrCQnJxcrRm9v72KdXx7l5SnK\n0qWKYm2tKO3aKcqZM1pHJIQoz4rzuVmi7XF37NjBmDFjaNKkCePHj+fixYtGPe/w4cO4urrSqlUr\nqlatysiRIwkLCzN4fkhICKNGjSpJiJXGnTvq9N3XX4enn1ZnY0lRASFEWTE6iZw9e5Z58+bRokUL\n/P39CQ4OpkePHnz77bdcunTJqNd4cL3Jw9aYpKenEx0dzfDhw3XHrKys8Pf3x9vbm6CgIGNDr7DO\nnlVXn2/ZAkuXqvuAyGxrIURZeuiYyM2bNwkJCeGzzz7jp59+QlEUunbtyvTp03njjTeYM2cOPXv2\nNPpiiqIUOmZloOZGeHg43bt3p8F904r279+Po6Mjf/zxB/3796dt27Z6rx8UFKRLMhV106ywMBgz\nRl3z8f330K+f1hEJISojgy2RESNG4OjoyLRp07h27Rpz5szhzJkzxMbGMmnSJL0JoShOTk4kJCTo\n7icmJuLo6Kj33NDQ0EJdWfnnOjg48NRTT3H48GG9zw0ICCAuLo64uDgaNWpU7DgtWW4u/Pvf6pa1\nrVur5UskgQghtGIwiWzZsoWsrCyeeOIJDh06xKJFi0pd+t3X15f4+HguXbpEVlYWoaGhDB06tNB5\nqamp7Nmzh2H3TS+6e/cud+7c0f35+++/x9PTs1TxlDc3bsDgwfDuuzB+POzbJ/t/CCG0ZTCJLFy4\nEFdXVyIjI2nRogWDBw/mq6++IjMzs8QXs7W1ZfXq1QwYMAB3d3eeffZZPDw8WLNmTYHtd7du3Yq/\nvz+1atXSHbt+/TqPP/44HTt2pEuXLgwaNIiBAweWOJbyIDgYXFzA2hocHdUB8x07YM0aWLsWqlfX\nOkIhRGVnpRTRL3XgwAE2btzI5s2buXXrFnXr1mXQoEGEhISwe/duevToUVaxloiPj0+5XE8SHAwB\nAZCefu+YlRW89ZZ6E0IIcynO52aRSSRfZmYm27ZtY+PGjcTExJCbm4uLiwsTJkxg7NixBqv8aq28\nJhEXF7hypfBxZ2e4fLmsoxFCVCZmSSL3u379Op9//jmff/45J06cwMbGhuzs7GIHWhbKaxKxtlaL\nKD7IykrdG10IIcylOJ+bJVps2LhxY2bPns3x48eJi4vj1VdfLcnLiIdo0aJ4x4UQQgslSiL369y5\nMytXrjRFLOI+ixYV3ra2Zk31uBBCWIpSJxFhHr16qd1Z9eqpycTZGYKC1O1thRDCUhhdxVeUrW+/\nVX/++CNUsuUwQohyRFoiFioyEpo3Bw8PrSMRQgjDJIlYoMxM+OEHGDSo8LiIEEJYEkkiFmjvXrh7\nF554QutIhBDi4YweE7l48SJfffUVV69eJSMjo8BjVlZWrF271uTBVVZRUVCtGvTtq3UkQgjxcEYl\nkbCwMJ555hny8vJwcHCgWrVqBR43VM5dlExkJPTuDfeVDhNCCItkVBJ588036d27N8HBwRWutLql\niY9Xb9OmaR2JEEIUzagkcvHiRZYvXy4JpAxERak/Bw3SNg4hhDCGUQPrbdu2JSUlxdyxCNQk4uYG\nrVppHYkQQhTNqCSydOlSFi9ezMWLF80dT6WWlga7d0srRAhRfhjVnRUYGEhKSgru7u60bt26wL7n\noA6s79mzxywBViY7d0JWlkztFUKUH0YlERsbG9zc3MwdS6UXGQm1a4OF7/MlhBA6RiWR3bt3mzkM\noSjqeEj//lC1qtbRCCGEcWTFuoU4cQISE2U8RAhRvhidRK5du8bs2bPx9fXlkUceoUuXLrz++uv8\n/vvv5oyv0sif2vuPf2gbhxBCFIdRSeTcuXN4eXnx4YcfUrt2bbp06UKtWrVYtWoVXl5exMfHmzvO\nCi8yEjp1AkdHrSMRQgjjGTUm8sYbb2BnZ8ehQ4dwcXHRHb9y5Qr+/v688cYbbNmyxVwxVng3b8KB\nAzB3rtaRCCFE8RjVEtm1axfvvPNOgQQC4OzsTGBgILt27TJHbJXGd99BXp6Mhwghyh+jkkhWVhZ1\n6tTR+1idOnXIysoyaVCVTVQUNGwIXbpoHYkQQhSPUUnEy8uLjz76iLy8vALHFUXhv//9L15eXmYJ\nrjLIzVW3wh04EGxstI5GCCGKx6gxkQULFjB48GDc3d157rnnaNq0Kb///jubN28mPj6eyMhIc8dZ\nYcXFwZ9/yip1IUT5ZFRLZODAgURERFCnTh0WLVrEq6++yrvvvkvt2rWJiIjA39/f6AtGR0fj5uaG\nq6srS5YsKfT4smXL8PLywsu2YgJtAAAgAElEQVTLC09PT2xsbLhx44bu8dzcXDp16sTgwYONvqYl\ni4wEa2sYMEDrSIQQovisFEVRivOE9PR0bt68Sf369alZs2axLpabm0ubNm2IiYnByckJX19fQkJC\naNeund7zw8PDWbFiBTt37tQd++CDD4iLi+P27dtEREQUeU0fHx/i4uKKFWdZ8vFRdzHcv1/rSIQQ\nQlWcz81ir1ivWbMmzZo1K3YCATh8+DCurq60atWKqlWrMnLkSMLCwgyeHxISwqhRo3T3ExMTiYyM\nZOLEicW+tiW6dg1+/llmZQkhyi+DYyILFy5k4sSJODo6snDhwoe+iJWVFf/+97+LvFhSUhLNmzfX\n3XdycuLQoUN6z01PTyc6OprVq1frjs2cOZOlS5dy586dIq9VHkRHqz9lPEQIUV4ZTCKBgYEMHDgQ\nR0dHAgMDH/oixiYRfT1nhvZnDw8Pp3v37rqy8xERETg4OODt7V1kQcigoCCCgoIASE5OLjIurURG\nqivUO3bUOhIhhCgZg0nk/um8D07tLSknJycSEhJ09xMTE3E0UOcjNDS0QFfW/v372b59O1FRUWRk\nZHD79m1eeOEFvvjii0LPDQgIICAgAFD79ixRdjZ8/z089xwYyKNCCGHxjBoTuXr1KtnZ2Xofy8nJ\n4erVq0ZdzNfXl/j4eC5dukRWVhahoaEMHTq00Hmpqans2bOHYcOG6Y699957JCYmcvnyZUJDQ+nb\nt6/eBFJe7NsHd+7IeIgQonwzKom0bNmSo0eP6n3s+PHjtGzZ0qiL2drasnr1agYMGIC7uzvPPvss\nHh4erFmzhjVr1ujO27p1K/7+/tSqVcuo1y2PoqKgShXo10/rSIQQouSMmuJrbW3NwYMH6aKnLsfB\ngwfp0aOHwZaK1ix1im+7dtCsGcTEaB2JEEIUVJzPTYNjIrdu3SqwyC8pKYmLFy8WOOevv/5i48aN\nNGnSpIShVk6XLsHp0zBpktaRCCFE6RhMIqtWreLtt9/GysoKKysrRowYofc8RVF4++23zRZgRZS/\nAZWMhwghyjuDSeTJJ5/ExcUFRVGYMGECb775Jo888kiBc6pVq0a7du3o0KGD2QOtSKKi4JFHoHVr\nrSMRQojSMZhEOnbsSMe/FzBYWVkxePBgGjZsWGaBVVTp6bBzJwQEyNReIUT5Z1QV37Fjx5o7jkpj\n927IyJBV6kKIisGoJAJw8uRJ1q5dy9mzZ8nIyCjwmJWVFTt27DB5cBVRZCTUrAm9emkdiRBClJ5R\nSeTQoUP06tULFxcX4uPj6dChAzdv3uTq1as4OTnh6upq7jgrBEVRx0P69YPq1bWORgghSs+oxYbz\n5s3j6aef5tdff0VRFNauXcvly5f54YcfyM3N5c033zR3nBXC6dNw+bLMyhJCVBxGJZFffvmFF154\nQVcsMTc3F4C+ffvy5ptvMnfuXPNFWIHkT+39xz+0jUMIIUzFqCSSnZ1NrVq1sLa2pkGDBly7dk33\nmJubGydPnjRbgBVJZCS0bw8tWmgdiRBCmIZRSeSRRx4hKSkJgA4dOrBu3Try8vLIy8tj/fr1smLd\nCKmpatFFmZUlhKhIjBpYHzJkCLt37+b5559n3rx5DBo0CDs7O2xsbEhLS+PDDz80d5zlXkwM5OTI\neIgQomIxKoncvymVn58fBw8e5JtvviE9PZ2BAwfi7+9vrvgqjKgoqFcPunXTOhIhhDAdo9eJ3K9T\np0506tTJ1LFUWHl5ahIZMABsS/SOCyGEZTJqTOTgwYN89dVXeh/bvHmzwX3SheroUbh+XcZDhBAV\nj1FJZO7cufz66696Hzt9+rRM8S1CZKRaJ2vgQK0jEUII0zIqiRw/fpxHH31U72NdunThl19+MWlQ\nFU1UFPj6goOD1pEIIYRpGZVEMjIyyMvL0/tYbm4ud+/eNWlQFUlyMhw+LLOyhBAVk1FJxN3dne3b\nt+t9bPv27bi5uZk0qIokOlqtmSXjIUKIisiouUKTJ0/m5Zdfxs7OjkmTJuHk5ERSUhJBQUGsXbuW\n//73v+aOs9yKjITGjaFzZ60jEUII0zMqiUyaNImzZ8+yYsUKPvjgA91xKysrZs2aRUBAgNkCLM9y\ncuC77+DJJ8HaqDafEEKUL0avWnj//fd55ZVX+OGHH0hJScHe3h4/Pz9atWplzvjKtdhYuHVLxkOE\nEBVXsZa+PfLII4X2WReGRUWpiwv799c6EiGEMA+DSeTq1as0bdqUKlWqcPXq1SJfqIWUpi0kMhIe\nfxzq1tU6EiGEMA+DScTFxYWDBw/SpUsXXFxcdHuJGJK/x4hQJSTAiROwdKnWkQghhPkYTCLr16/X\ndV2tW7euyCQiCsrfgErGQ4QQFZnBJJKamlpgB8P8ri1hnKgocHYGd3etIxFCCPMxOPF01qxZXL58\nGYCWLVty9OhRk1wwOjoaNzc3XF1dWbJkSaHHly1bhpeXF15eXnh6emJjY8ONGzfIyMigS5cudOzY\nEQ8PD9566y2TxGMOGRnwww9qK0QacEKIisxgEqlXrx6///47AIqimKQ7Kzc3l1dffZVvv/2WU6dO\nERISwqlTpwqc869//Ytjx45x7Ngx3nvvPXr16kWDBg2oVq0aO3fu5Pjx4xw7dozo6GgOHjxY6pjM\nYe9eSE+XVepCiIrPYHdW9+7dGTt2LB07dgTglVdewc7OTu+5VlZW7Nixo8iLHT58GFdXV93akpEj\nRxIWFka7du30nh8SEsKoUaN016hduzag7vmenZ1tseM0kZFQvTr06aN1JEIIYV4GWyKffvopo0aN\nwtraGisrK3JycnQf3g/esrKyjLpYUlISzZs3193PL5+iT3p6OtHR0QwfPlx3LDc3Fy8vLxwcHOjf\nvz9du3bV+9ygoCB8fHzw8fEhOTnZqNhMKSpKTSA1a5b5pYUQokwZbIk0btxYVxPL2tqaoKAgunTp\nUqqLKYpS6Jih1kR4eDjdu3enQYMGumM2NjYcO3aMW7du8dRTT3Hy5Ek8PT0LPTcgIEBXisXHx6dU\nMRfXuXNw/jzMnFmmlxVCCE0YVdHp0qVLeHl5lfpiTk5OJCQk6O4nJibi6Oio99zQ0FBdV9aD6tWr\nR+/evYmOji51TKaWP7VXxkOEEJWBUUnE2dmZqlWrlvpivr6+xMfHc+nSJbKysggNDWXo0KGFzktN\nTWXPnj0MGzZMdyw5OZlbt24B8Ndff/HDDz/Qtm3bUsdkapGR6rTeli21jkQIIczPYHeWjY0NsbGx\ndOnSRTcuYkj+mEmRF7O1ZfXq1QwYMIDc3FwmTJiAh4cHa9asAdSS8wBbt27F39+fWrVq6Z577do1\nxo4dS25uLnl5eTz77LMMHjzY6F+0LKSlwZ49MGOG1pEIIUTZsFL0DVQAb7/9NpMmTcLR0ZHAwMAi\nZ0JZ6roNHx8f4uLiyuRa27bBU0/Bzp0yM0sIUX4V53PTYBKpKMoyiQQEQGgopKSALO4XQpRXxfnc\nLPFWSTdu3ODnn38mMzOzpC9RoSiKOqju7y8JRAhReRiVRN59913mzp2ru793715cXFzo0qULrVu3\nJj4+3mwBlhe//AJJSVJwUQhRuRiVRL744osCOxi+/vrrdOzYkW3bttG4cWP+/e9/my3A8iIyUv35\nj39oG4cQQpQlo3Y2TEpKonXr1oA61fann35ix44d9O7dm6ysLKZPn27WIMuDqCjw9oYmTbSORAgh\nyo5RLREbGxtdaZO9e/dSvXp1unfvDkCjRo24ceOG+SIsB1JS1P3UZYGhEKKyMSqJeHh48MUXX5CW\nlsa6devo1auXbm+RhIQEHBwczBqkpfv+e8jLk/EQIUTlY1R31oIFCxg2bBjBwcFUqVKF7777TvdY\nVFQUnTt3NluA5UFkJNjbQxmX6RJCCM0ZlUQGDBjA6dOnOXLkCF5eXrptcwF69uypKxdfGeXmQnS0\n2pVlY6N1NEIIUbaMSiKg7m7YUk9BqJdfftmkAZU3hw+rYyIyHiKEqIyMGhMJCwtj/fr1uvtXrlyh\nW7du1KlThxEjRpCWlma2AC1dVBRYW8OAAVpHIoQQZc/oxYb3b+702muvkZiYSEBAAHv37iUwMNBc\n8Vm8yEh47DGoX1/rSIQQouwZlUQuXLhAhw4dALUMe1RUFB988AHLly9n8eLFbN261axBWqrffoOj\nR2VWlhCi8jIqiWRkZFCjRg0ADhw4QE5ODv7+/gC4ubnx22+/mS9CC/btt+pPGQ8RQlRWRiURFxcX\n9u3bB6jjI97e3tStWxeAP/74Q/fnyiYqCpycoH17rSMRZSY4GFxc1IEwFxf1vhCVmFGzs15++WVm\nz57N1q1bOXbsGB9//LHusdjYWNq1a2e2AC1VVhbExMCoUVDEViuioggOVuv9p6er969cUe8DjB6t\nXVxCaMioJDJjxgzs7e05ePAg06dPZ8yYMbrH7ty5w/jx480WoKXatw/u3JHxkEpl/vx7CSRferp6\nXJKIqKRkUyp9goPVD4arV6FFC1i0qNCHxD//CatXq2tEatc2YcDC8mRlwd690L+/4XM8PdW+zfxb\ns2YF79etK01WUW4U53PT6MWGlYaRXRZRUdC7tySQCuvPP9W/5PBw+O47tdlpSJ064OoKiYlw7Bhc\nv67uUna/WrUKJhV9icbeXhKNKHeMTiJBQUF8/PHHnD17Vu9uhrm5uSYNTDNGdFlcvAhnzsDkyRrE\nJ8xDUeD0aTVphIerZZnz8qBpUxg5EoYMURPL1KkF/33UrAkff1ywpZqVBdeuqUkl/5aUdO/PO3eq\n88Mf/D9TrVrhxPJgsmncWOrrCItiVBL57LPPmDZtGmPHjuX48eNMmDCB7Oxstm/fTqNGjRhdkfqD\nr14t8nhUlPpTxkPKufxuqvBwiIhQvx0AdOoEb76pJo7OndWZWPmqVi2yq5OqVcHZWb0ZkpurtlgM\nJZqDB9Wff2/BoGNjA46OD2/RODrKHs2izBg1JtK5c2eGDh3Kv//9b6pUqUJcXBydO3fm5s2b9O7d\nm0mTJjF16tSyiLfYij0m4uKidmE9qFEj+OMPQN298MIFOHfONDGKMpSSUrCb6vZtqF4d+vVTk8ag\nQeoHsSVQFLX1c3+ieTDZJCQUbjlbWaktloe1aJo1g7/XfgnxIJOPicTHx9OzZ0+sra2xtrbWbVBV\nv3595s+fz/z58y02iRTbokUFx0RA/SaanAxLl3J3yr/YtcuKV17RLkRRDPndVBERauI4cEDtpmrS\nBJ59Vk0c/fqpYxaWxspK/fLSqJHaOtJHUSA11XCiOX8edu+GW7cKP7dhw6LHaerUMeuvKMo/o5JI\njRo1yMvLw8rKiiZNmnDx4kUeffRRAGrXrl2xVqznd03c32Xx1lvqt9Y33mDX91XJzJwpq9QtWVYW\n/PjjvfGN+7up5s9XE4e3d8FuqvLKygrq1VNvnp6Gz0tLK9iCebBFc/iw+kXpQXZ2RSea+vVlQkAl\nZlQSad++PefPn8fPz48ePXqwePFiWrZsia2tLYGBgbRt29bccZat0aML93OPGwdt2xL5djVqWafT\ns106YK9FdEKflBS1Dk14uLrBy+3b6kB1v37wr3/B4MGW002lhdq1wc1NvRmSkaEO+Btq1Zw8qU4Y\neLAHvEaNwonmwWTTqFHFSNqiEKOSSEBAABf//jb3zjvv4Ofnx+OPPw5AnTp12LZtm/kitBRWVihv\nBRL1n7v43dhBtR4z1C6SSrha3yIoijpFLn9QfP9+tZuqcWN45hm1teHnZ5ndVJaqenVo1Uq9GZKd\nDb//bjjR7N2r/jknp+DzqlQpeuZZkyZgK6sOypsSLTa8e/cusbGxpKen89hjj2Fvb/w38ujoaGbM\nmEFubi4TJ05kzpw5BR5ftmwZwX/XI8rJyeH06dMkJydz9+5dxowZw++//461tTUBAQHMmDGjyOuV\naLGhASdPqnWyguZeYtL6x9Rxky+/hIEDTfL6ogjZ2QW7qS5cUI97ealJoyJ1U5VneXnqJJQHE82D\nXWgZGQWfZ22tTql+WKJxdFRbmMKsivO5WaYr1nNzc2nTpg0xMTE4OTnh6+tLSEiIwdpb4eHhrFix\ngp07d3Lt2jWuXbtG586duXPnDt7e3mzbtq3Iul2mTCJLl8Ibb6gTYpyUBBg6FH75BVasgGnTpF/Y\nHPK7qSIi1G6q1FT1Q6RvXzVpDB4MzZtrHaUoLkWBGzcenmgSEtSxnAc5OBQ9TlOzZtn/ThWISWZn\nXTW0XsKAFi1aFHnO4cOHcXV1pdXfzeWRI0cSFhZmMBGEhIQwatQoAJo2bUrTpk0BtQvN3d2dpKSk\nMi3+GBkJHTvmd603V78Vv/gizJgBp07BRx/J/PzSUhQ4e/Zea+P+bqrhw9XE0b+/dFOVd1ZW6uyw\nhg3V/1SG3L5tONFcvqwWsbtxo/Dz6tcvepzGzk6++JmAwSTi4uKCVTHeYGNWrCclJdH8vm+NTk5O\nHDp0SO+56enpREdHs3r16kKPXb58maNHj9K1a1ej4yutW7fUz7PXX7/vYO3a8M036oyfJUvU6ZSb\nN8s2h8WVna1+GOQnjvPn1eMdO8K8eWri8PGRbqrKyM5OHXd82JfF9PTCM8/uTzZHjqgLOx9Uu3bR\nLZqGDSXRFMFgElm3bl2xkogx9PWcGbpGeHg43bt3p0GDBgWOp6WlMXz4cFauXImdnZ3e5wYFBREU\nFARQYFvf0vj+e3WRcaFV6tbW8N574O4OkybBo4+qH4Rt2pjkuhXWjRv3uqm+/VbtpqpaVe2mmjVL\n7aYyonUrBDVrQuvW6s2QzMzCpWjuTzQ//KDOTMvLK/i8atWKbtE4OFTqUjQGk8i4ceNMfjEnJycS\nEhJ09xMTE3F0dNR7bmhoqK4rK192djbDhw9n9OjRPP300wavExAQQMDfRRN9fHxMELm6yLl+fTDY\n+BkzRp3V8tRTaiL5+mv1A1Hc82A3VW6u+h8wv5vKz08qWgrzqFZNrUbh4mL4nJycwqVo7r8dOKD+\nzM4u+Dxb28KlaB5MNE2bVtiuboMD64qiEBERQcuWLfE0sIjpxIkTXL58mSFDhhh1sZycHNq0acOO\nHTto1qwZvr6+bNq0CQ8PjwLnpaam0rJlSxISEqj1d9+3oiiMHTuWBg0asHLlSqN/QVMMrOfX4evX\nDzZtKuLkS5fUD8SzZ9Va8S+/XKprl2v53VT5q8Xj49XjHTrcm03l6yvdVKL8yMvTX4rm/lZNQgL8\n9VfB51lZqVOYiypFU726Nr/XA0wysP75558zZcoUTpw4YfDJderUYdSoUXz66aeFWg16L2Zry+rV\nqxkwYAC5ublMmDABDw8P1qxZA8Dkv8vibt26FX9/f10CAdi/fz+ff/457du3x8vLC4DFixfzRBks\nHf/5Z3XGolGXatlS/cYyapRa5vf0aXj//coz//3mzYKL/m7dutdNNXOm2h/4sMKEQlgya2u19ezg\noBbn1EdR1H/3hhLNuXNqJefU1MLPtbcvuvusqNa6EfshmZLBloi/vz9ubm589NFHD32BGTNmcPbs\nWaKjo80SYGmZoiUSGAgLF6qJxOglMbm5MHs2rFypVmwMCVE3JqqIzp271021b9+9bqpBg+7NppJu\nKiEKunPn4RMCEhPVVs+D6tY1nGiOH1c/rO5vCdWsCUFBxUokJlknYm9vz/r164vsqgoPD2f8+PH8\nqe+XtQCmSCJduqjjZrGxJXhyUBC8+qo60B4e/vDVwOVFTk7B2VT3d1MNHqwmji5dpJtKiNL66y/9\npWjuTza//164FM2DnJ3VKdFGMkl31p07d6hvxFTV+vXrc+dhu76Vc9evw08/wTvvlPAFAgLUWSPD\nh6uj8lu2QI8eJo2xTNy8qXZPhYer3VX53VR9+qjrZAYPlm4qIUytRg145BH1Zkh29r2ZZ48/rj+h\nFHPdX3EYTCL29vZcuXJFVyPLkKtXrxar7El5k99LV6qhlz594NAh9YO2Xz+1dWKG2W8ml99NFRGh\nLqzMzVUL6T355L1uKikVLoS2qlRRxz7yb/r2QzLjdHmD/Q2PP/44GzduLPIFNmzYUGSiKc8iI9WZ\nWYa2czBa69bqbnW9esH48Wr9FEvbUjgnB/bsUcdy8iu+zp6tlh554w21P+/aNVi/Hp5+WhKIEJZm\n0aLCJV9q1lSPm4tiwIEDBxRra2tl5syZSmZmZqHHs7KylOnTpys2NjZKbGysoZfRnLe3d4mfm5Wl\nKHXrKsqECSYMKCtLUV55RVFAUYYOVZQ7d0z44iVw86aihIQoyvPPK0r9+mpcVasqir+/onz0kaJc\nuqRtfEKI4vniC0VxdlYUKyv15xdfFPslivO5abA7q1u3bixfvpx//vOfBAcH4+/vj/Pffd5Xrlwh\nJiaGlJQUli9frtugqqI5cECdhWfSvdSrVIH//lct4zBjhtqHuX172a7Ojo+/Nyie301lb68WlBwy\nBPz9pZUhRHmlbz8kM3ro4oWZM2fSuXNnlixZwtatW/nr72ljNWrUoHfv3syZM4ce5XGQ2EhRUepn\nvp+fGV586lS1i+vZZ9WZTNu2qSvdzSEnR82I+Ynj7Fn1uKenWgwsfzZVJS7dIIQoGaNLwefl5emm\n8TZs2BCbcvKBU5opvp6eavHYHTtMHNT9Tp9WB9yTktSxBiMWbRrl1i11VkBEhJoNb95UM2Lv3vdW\niz+sBIQQotIyyRTfB1lbW+Pg4FDioMqT4GB1HDkpSb0FB5uxdejurs7cGj4cnn9eTSqBgSVbY3H+\nfMFuqpycgt1U/furVVGFEMJEKkktDuMFB6tLO9LT1fu3bqn3wYyJxN4eYmLUMinvvKNu+7phQ9Eb\n6+R3U+XXpjpzRj3u4aHOqhoyRF2bUk5ajUKI8keSyAPmz7+XQPKlp6vHzTpWVbUqrF2rDri//rpa\nyHHsWLXu1v01cAYPLrjo78aNe91UU6aoj7dsacZAhRDinjLdHlcLxR0TsbbWv+DTyqrwVgNmEx4O\nzzwDWVkFg8kPTlHU1ssTT9ybTSXdVEIIEzHLmEhlocGCz8KGDIEGDdSFfffLy1OTRVSUOpNLuqmE\nEBqTCnkP0GLBp16//67/+J070L27JBAhhEWQJPKA0aPV0lbOzmoXlrNzsasom4ahpo9sGSuEsCCS\nRPQYPVqtmpyXp/4s8wQCFtQkEkIIwySJWCqLaRIJIYRhMrBuycq4Bo4QQhSXtESEEEKUmCQRIYQQ\nJSZJRAghRIlJEhFCCFFikkSEEEKUWIWvnWVvb49LOd83Izk5mUaNGmkdhkWQ96IgeT8KkvfjntK8\nF5cvX9btH1WUCp9EKoLSbKxV0ch7UZC8HwXJ+3FPWb0X0p0lhBCixCSJCCGEKDGbwMDAQK2DEEXz\n9vbWOgSLIe9FQfJ+FCTvxz1l8V7ImIgQQogSk+4sIYQQJSZJxEIlJCTQp08f3N3d8fDwYNWqVVqH\npLnc3Fw6derE4MGDtQ5Fc7du3WLEiBG0bdsWd3d3YmNjtQ5JUytWrMDDwwNPT09GjRpFRkaG1iGV\nqQkTJuDg4ICnp6fu2I0bN+jfvz+tW7emf//+3Lx50yzXliRioWxtbVm+fDmnT5/m4MGD/Oc//+HU\nqVNah6WpVatW4e7urnUYFmHGjBkMHDiQM2fOcPz48Ur9viQlJfHhhx8SFxfHyZMnyc3NJTQ0VOuw\nytS4ceOIjo4ucGzJkiX069eP+Ph4+vXrx5IlS8xybUkiFqpp06Z07twZgDp16uDu7k5SUpLGUWkn\nMTGRyMhIJk6cqHUomrt9+zZ79+7lpZdeAqBq1arUq1dP46i0lZOTw19//UVOTg7p6ek4OjpqHVKZ\n6tmzJw0aNChwLCwsjLFjxwIwduxYtm3bZpZrSxIpBy5fvszRo0fp2rWr1qFoZubMmSxduhRra/kn\ne/HiRRo1asT48ePp1KkTEydO5O7du1qHpZlmzZoxe/ZsWrRoQdOmTalbty7+/v5ah6W569ev07Rp\nU0D9UvrHH3+Y5TryP9LCpaWlMXz4cFauXImdnZ3W4WgiIiICBwcHmbr5t5ycHI4cOcIrr7zC0aNH\nqVWrltm6KsqDmzdvEhYWxqVLl/jtt9+4e/cuX3zxhdZhVRqSRCxYdnY2w4cPZ/To0Tz99NNah6OZ\n/fv3s337dlxcXBg5ciQ7d+7khRde0DoszTg5OeHk5KRrmY4YMYIjR45oHJV2fvjhB1q2bEmjRo2o\nUqUKTz/9NAcOHNA6LM01btyYa9euAXDt2jUcHBzMch1JIhZKURReeukl3N3dee2117QOR1Pvvfce\niYmJXL58mdDQUPr27Vupv2k2adKE5s2bc/bsWQB27NhBu3btNI5KOy1atODgwYOkp6ejKAo7duyo\n1BMN8g0dOpSNGzcCsHHjRoYNG2aW60gSsVD79+/n888/Z+fOnXh5eeHl5UVUVJTWYQkL8dFHHzF6\n9Gg6dOjAsWPHmDdvntYhaaZr166MGDGCzp070759e/Ly8ggICNA6rDI1atQounXrxtmzZ3FycmLt\n2rXMmTOHmJgYWrduTUxMDHPmzDHLtWXFuhBCiBKTlogQQogSkyQihBCixCSJCCGEKDFJIkIIIUpM\nkogQQogSkyQiRBnZsGEDVlZWnD9/vsDxn376iQYNGtCpUyf+/PNPjaITomQkiQihoQMHDuDn50fr\n1q3ZuXMn9vb2WockRLFIEhFCI3v27GHAgAG0b9+emJgY6tevr3VIQhSbJBEhNBATE8M//vEPfH19\n+e677yptcU1R/kkSEaKMRUZGMmTIEHr27ElkZCS1atXSOiQhSkySiBBlbObMmTg5OREWFkaNGjW0\nDkeIUpEkIkQZGzRoEBcuXOC9997TOhQhSs1W6wCEqGxWrFhBkyZNePvtt6levbrZqqsKURYkiQhR\nxqysrAgKCiIzM5O5c+dSvXp1Zs6cqXVYQpSIJBEhNGBtbc2GDRvIyspi1qxZVK9encmTJ2sdlhDF\nJklECI3Y2NgQHBxMZv7q8nMAAAB4SURBVGYmU6ZMoVq1aowfP17rsIQoFhlYF0JDtra2fPXVVwwc\nOJCJEyeyadMmrUMSolhkZ0MhhBAlJi0RIYQQJSZJRAghRIlJEhFCCFFikkSEEEKUmCQRIYQQJSZJ\nRAghRIlJEhFCCFFikkSEEEKUmCQRIYQQJfb/CeiSWyLE22EAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {
      "tags": []
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "encoded_size = 2\n",
    "num_hidden = 128\n",
    "nz = 1\n",
    "beta = 0.05\n",
    "lr = 0.001\n",
    "\n",
    "K = [1, 2, 5, 10]\n",
    "objectives = [Selbo, Siwae]\n",
    "\n",
    "acc_selbo = []\n",
    "acc_siwae = []\n",
    "\n",
    "for k in K:\n",
    "  for objective in objectives:\n",
    "    # Setup the model\n",
    "    loss_fn = objective(\n",
    "        encoder=Encoder(input_size=info.features['image'].shape[0],\n",
    "                        encoded_size=encoded_size,\n",
    "                        num_hidden=num_hidden,\n",
    "                        nmix=k),\n",
    "        decoder=Decoder(encoded_size=encoded_size,\n",
    "                        num_classes=info.features['label'].num_classes),\n",
    "        prior=tfd.MultivariateNormalDiag(loc=tf.zeros(encoded_size),\n",
    "                                        scale_diag=tf.ones(encoded_size)),\n",
    "        nz=nz,\n",
    "        beta=beta,\n",
    "        seed=seedstream)\n",
    "\n",
    "    opt = tf.optimizers.Adam(lr)\n",
    "    train_op = tfn.util.make_fit_op(\n",
    "        loss_fn=loss_fn, optimizer=opt,\n",
    "        trainable_variables=loss_fn.trainable_variables,\n",
    "        xla_compile=False)  # xla doesn't like this... not sure why\n",
    "\n",
    "    # Train the model\n",
    "    for i, (x, y) in enumerate(iter(train_data)):\n",
    "      this_loss, (this_nll, this_kl) = train_op(x, y)\n",
    "\n",
    "    if objective == Selbo:\n",
    "      acc_selbo.append(evaluate_accuracy(\n",
    "          encoder=loss_fn.encoder,\n",
    "          decoder=loss_fn.decoder,\n",
    "          dataset=eval_data,\n",
    "          nz=10000,\n",
    "          seed=seedstream).numpy())\n",
    "    else:\n",
    "      acc_siwae.append(evaluate_accuracy(\n",
    "          encoder=loss_fn.encoder,\n",
    "          decoder=loss_fn.decoder,\n",
    "          dataset=eval_data,\n",
    "          nz=10000,\n",
    "          seed=seedstream).numpy())\n",
    "\n",
    "plt.plot(K, acc_selbo, 'r-', label='Selbo (nz=1)')\n",
    "plt.plot(K, acc_selbo, 'ro')\n",
    "plt.plot(K, acc_siwae, 'b-', label='Siwae (nz=1)')\n",
    "plt.plot(K, acc_siwae, 'bo')\n",
    "plt.xlabel('K', fontsize=16)\n",
    "plt.ylabel('Classification Accuracy', fontsize=16)\n",
    "plt.legend()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "single_column_mnist.ipynb",
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
